try:
  import pylab
except ImportError:
  print 'Could not import pylab.'
import sys

def file_out (filename, dist, coll):
  """
  Output the average distance as a funtion of time to a specified file
  """
  dist_av = dist.crop_av_distance (dist.dist_av, dist.w)
  wall_coll = coll.crop_av_collisions (coll.wall_coll, coll.w)
  bot_coll = coll.crop_av_collisions (coll.bot_coll, coll.w)
  # note bot and wall collisions same by definition
  if wall_coll and bot_coll and dist_av:
    if len (wall_coll) != len (dist_av):
      print 'Error: wall_coll and dist_av not same length. They should be.'
    if len (bot_coll) != len (dist_av):
      print 'Error: bot_coll and dist_av not same length. They should be.'
    file = open (filename, 'w')
    str_data = []
    file.write ('# total distance :: total wall colls :: total bot colls')
    for i in xrange (0, len(dist_av)):
      try:
        str_data.append (str(dist_av[i])+'\t'+str(wall_coll[i])+'\t'+str(bot_coll[i])+'\n')
      except IndexError:
        print 'Index out of range....'
    file.writelines (str_data)
    file.close
    
class distance_tracker:

  def __init__ (self):
    # i^th entry is t = i...;
    self.dist_av = []
    # Corresponding weights
    self.w = []
  
  def av_dist (self, tot_dist, t, t_trav, v_trav):
    ti = int (t - t_trav+1) # Need the +1 in order to not over-extrapolate
    tf = int (t)
    df = tot_dist # Final total distance
    if len (self.w) < tf + 1:
      for i in xrange (len(self.w), tf+1):
        # make w the same length as time
        self.w.append (0)
    if len (self.dist_av) < tf + 1:
      for i in xrange (len(self.dist_av), tf+1):
        # Make the av_dist the same lenght as the int. number of time steps
        self.dist_av.append (0)
  
    # Note these list lengths shouldn't always be as long as tf, if they
    # were than they would have to be reset each time a trial started (henc
    # new event occurring before tf...). However, they should be the same
    # lenght as each other.
    if len (self.dist_av) != len (self.w):
      print 'Warning: av_dist and weights not the same length...'

    for i in xrange (ti, tf+1):
      new_dist = df - v_trav*(tf-i)
      self.dist_av[i] = (self.w[i]*self.dist_av[i]+new_dist)/(self.w[i]+1)
      self.w[i] += 1
  
  def crop_av_distance (self, dist_av, w):
    """
    Chop off the tails where not all trials make it
    """
    # Find the max weight. This corresponds to the steps that all trials
    # contributed to (i.e. the beginning/middle ones)
    if w:
      mw = max(w)
      for i in xrange (0, len(w)):
        if w[i] < mw:
          # For each step w/o good weights, we must remove that spot.
          # If N spots have bad weights, we remove the last N spots
          dist_av.pop()
      return dist_av
    else:
      return 0

  def graph_av_distance (self):
    """
    Show the average distance as a function of time in a pylab plot
    """
    dist_av = self.crop_av_distance (self.dist_av, self.w)
    pylab.title ('Bot distance vs. Time')
    pylab.xlabel ('Time')
    pylab.ylabel ('Total Distance per Bot')
    pylab.plot (dist_av)
    pylab.grid ()
    pylab.show ()

class collision_tracker:
  def __init__ (self):
    # i^th entry is t = i...;
    self.bot_coll = []
    self.wall_coll = []
    # Corresponding weights
    self.w = []
    
  def av_coll (self, tot_wall_coll, tot_bot_coll, t, t_trav):
    """
    Inputs: tot_wall_coll, tot_bot_coll, t, t_trav
    """
    ti = int (t - t_trav+1) # Need the +1 in order to not over-extrapolate
    tf = int (t)
    
    if len (self.w) < tf + 1:
      for i in xrange (len(self.w), tf+1):
        # make w the same length as time
        self.w.append (0)
    if len (self.bot_coll) < tf + 1:
      for i in xrange (len(self.bot_coll), tf+1):
        # Make the av_coll the same lenght as the int. number of time steps
        self.bot_coll.append (0)
    if len (self.wall_coll) < tf + 1:
      for i in xrange (len(self.wall_coll), tf+1):
        # Make the av_coll the same lenght as the int. number of time steps
        self.wall_coll.append (0)
  
    if len (self.wall_coll) != len (self.w):
      print 'Warning: av_coll and weights not the same length...'

    # Add last collision amount to all up to tf. That is, add the last
    # total collisions up to all but the current. Next time around we add
    # the current one, thus we only need to track the last total coll
    for i in xrange (ti, tf+1):
      self.wall_coll[i] = (self.w[i]*self.wall_coll[i]+float(tot_wall_coll))/(self.w[i]+1.)
      self.bot_coll[i] = (self.w[i]*self.bot_coll[i]+float(tot_bot_coll))/(self.w[i]+1.)
      self.w[i] += 1
  
  def crop_av_collisions (self, coll_av, w):
    """
    Chop off the tails where not all trials make it
    """
    # Find the max weight. This corresponds to the steps that all trials
    # contributed to (i.e. the beginning/middle ones)
    if w:
      mw = max(w)
      for i in xrange (0, len(w)):
        if w[i] < mw:
          # For each step w/o good weights, we must remove that spot.
          # If N spots have bad weights, we remove the last N spots
          coll_av.pop()
      return coll_av
    else:
      return 0

  def graph_colls (self):
    """
    Show the average distance as a function of time in a pylab plot
    """
    wall_coll = self.crop_av_collisions (self.wall_coll, self.w)
    bot_coll = self.crop_av_collisions (self.bot_coll, self.w)
    pylab.title ('Bot distance vs. Time')
    pylab.xlabel ('Time')
    pylab.ylabel ('Total Collisions per Bot')
    wall_graph = pylab.plot (wall_coll)
    bot_graph = pylab.plot (bot_coll)
    pylab.legend ((wall_graph,bot_graph),('Wall Collisions','Bot Collisions'),\
                  loc='upper left')
    pylab.grid ()
    pylab.show ()

class policy_tracker:
  def __init__ (self, n_states, n_actions):
    self.av_policy = []
    for i in xrange(0, n_states):
      self.av_policy.append ([])
      for j in xrange(0, n_actions):
        self.av_policy[i].append (0)
    self.n_states = n_states
    self.n_actions = n_actions
    self.w = 0

  def update_av_policy (self, bots):
    for bot in bots:
      for i in xrange (0, self.n_states):
        for j in xrange (0, self.n_actions):
          self.av_policy[i][j] = (self.w*self.av_policy[i][j] + bot.bra.policy[i][j])/(self.w + 1.)
      self.w += 1

  def write_policy (self):
    pol_file = open ('pol', 'w')
    for i in xrange (0, self.n_states):
      for j in xrange (0, self.n_actions):
        pol_file.write(str(i)+'\t')
        pol_file.write(str(j)+'\t')
        pol_file.write(str(self.av_policy[i][j])+'\n')
      pol_file.write('\n')
      



 
